// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#include <cmath>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>

#include <fmt/printf.h>
#include <glog/logging.h>
#include <pybind11/functional.h>
#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include "coin1/base/exception_util.h"
#include "coin1/fastfeed/feed_cache.h"
#include "coin1/fastfeed/proto_feed_builder.h"
#include "coin1/fastfeed/view/book_builder.h"
#include "coin1/fastfeed/view/string_json_builder.h"
#include "coin1/fastfeed/view/trade_builder.h"

namespace fastfeed {

namespace py = pybind11;

struct PyBookLevel {
  double price = 0;
  double qty = 0;
  int64_t count = 0;

  template <typename T>
  static PyBookLevel FromBookLevel(const T* param, const view::BookLevel& level) {
    return PyBookLevel {
        level.price / param->price_multiplier(),
        level.qty / param->qty_multiplier(),
        level.count};
  }

  template <typename T>
  static std::optional<PyBookLevel> FromOptionalBookLevel(
      const T* param,
      std::optional<view::BookLevel> level) {
    if (!level.has_value()) {
      return std::nullopt;
    }
    return PyBookLevel {
        level->price / param->price_multiplier(),
        level->qty / param->qty_multiplier(),
        level->count};
  }

  std::string to_str() {
    if (count == 0) {
      return fmt::format("BookLevel(price={}, qty={})", price, qty);
    } else {
      return fmt::format("BookLevel(price={}, qty={}, count={})", price, qty, count);
    }
  }
};  // struct PyBookLevel

class PyBookIterator {
 public:
  PyBookIterator(
      const view::detail::BookBuilderParameterHolder* param,
      view::BookIterator iter,
      view::BookIterator end)
      : param_(param), iter_(iter), end_(end) {}
  PyBookIterator(const PyBookIterator& other)
      : param_(other.param_), iter_(other.iter_), end_(other.end_) {}

  PyBookLevel next() {
    if (iter_ == end_) {
      throw py::stop_iteration();
    }
    return PyBookLevel::FromBookLevel(param_, *(iter_++));
  }

 private:
  const view::detail::BookBuilderParameterHolder* param_;
  view::BookIterator iter_;
  view::BookIterator end_;
};  // class PyBookIterator

py::array_t<double> levels_as_array(
    int num_level,
    view::BookIterator iter,
    view::BookIterator end,
    double price_multiplier,
    double qty_multiplier) {
  py::array_t<double> ret({num_level, 2});
  for (int i = 0; i < num_level; ++i) {
    if (iter != end) {
      *ret.mutable_data(i, 0) = (*iter).price / price_multiplier;
      *ret.mutable_data(i, 1) = (*iter).qty / qty_multiplier;
      ++iter;
    } else {
      *ret.mutable_data(i, 0) = 0;
      *ret.mutable_data(i, 1) = 0;
    }
  }
  return ret;
}

class PyBookView {
 public:
  explicit PyBookView(const view::BookBuilder* view) : view_(view) {}
  ~PyBookView() = default;

  const std::string& symbol() const { return view_->symbol(); }
  int64_t timestamp() const { return view_->timestamp(); }
  int64_t origin_timestamp() const { return view_->origin_timestamp(); }

  bool has_ask() const { return !view_->IsAskEmpty(); }

  bool has_bid() const { return !view_->IsBidEmpty(); }

  int32_t depth_ask() const { return view_->GetAskDepth(); }

  int32_t depth_bid() const { return view_->GetBidDepth(); }

  std::optional<PyBookLevel> ask0() const {
    return PyBookLevel::FromOptionalBookLevel(view_, view_->GetAsk0());
  }
  std::optional<PyBookLevel> bid0() const {
    return PyBookLevel::FromOptionalBookLevel(view_, view_->GetBid0());
  }

  py::array_t<double> get_bid_array(int num_level) {
    return levels_as_array(
        num_level,
        view_->BidBegin(),
        view_->BidEnd(),
        view_->price_multiplier(),
        view_->qty_multiplier());
  }

  py::array_t<double> get_ask_array(int num_level) {
    return levels_as_array(
        num_level,
        view_->AskBegin(),
        view_->AskEnd(),
        view_->price_multiplier(),
        view_->qty_multiplier());
  }

  PyBookIterator iter_ask() const {
    return PyBookIterator(view_, view_->AskBegin(), view_->AskEnd());
  }

  PyBookIterator iter_bid() const {
    return PyBookIterator(view_, view_->BidBegin(), view_->BidEnd());
  }

  // Returns sum of qty in [from_price, to_price) from ask book.
  double get_accumulated_ask_qty(double from_price, double to_price) {  // [from_price, to_price)
    if (from_price > to_price) {
      std::swap(from_price, to_price);
    }
    int64_t multd_from_price = std::llround(from_price * view_->price_multiplier());
    int64_t multd_to_price = std::llround(to_price * view_->price_multiplier());
    int64_t accumulated_qty = 0;
    auto it = view_->FindAskLevelByPrice(multd_from_price, true);
    for (; it != view_->AskEnd() && it->price < multd_to_price; ++it) {
      accumulated_qty += it->qty;
    }
    return accumulated_qty / view_->qty_multiplier();
  }

  // Returns sum of qty in [from_price, to_price) from bid book.
  double get_accumulated_bid_qty(double from_price, double to_price) {
    if (from_price < to_price) {
      std::swap(from_price, to_price);
    }
    int64_t multd_from_price = std::llround(from_price * view_->price_multiplier());
    int64_t multd_to_price = std::llround(to_price * view_->price_multiplier());
    int64_t accumulated_qty = 0;
    auto it = view_->FindBidLevelByPrice(multd_from_price, true);
    for (; it != view_->BidEnd() && it->price > multd_to_price; ++it) {
      accumulated_qty += it->qty;
    }
    return accumulated_qty / view_->qty_multiplier();
  }

  std::optional<double> get_ask_price_by_amount(double amount) {
    double acc_amount = 0.;
    for (auto it = view_->AskBegin(); it != view_->AskEnd(); ++it) {
      double price = it->price / view_->price_multiplier();
      double qty = it->qty / view_->qty_multiplier();
      acc_amount += price * qty;
      if (acc_amount >= amount) {
        return price;
      }
    }
    return std::nullopt;
  }

  std::optional<double> get_bid_price_by_amount(double amount) {
    double acc_amount = 0.;
    for (auto it = view_->BidBegin(); it != view_->BidEnd(); ++it) {
      double price = it->price / view_->price_multiplier();
      double qty = it->qty / view_->qty_multiplier();
      acc_amount += price * qty;
      if (acc_amount >= amount) {
        return price;
      }
    }
    return std::nullopt;
  }

  std::optional<double> get_ask_price_by_qty(double qty) {
    int64_t multd_qty = std::llround(qty * view_->qty_multiplier());
    int64_t acc_qty = 0;
    for (auto it = view_->AskBegin(); it != view_->AskEnd(); ++it) {
      acc_qty += it->qty;
      if (acc_qty >= multd_qty) {
        return it->price / view_->price_multiplier();
      }
    }
    return std::nullopt;
  }

  std::optional<double> get_bid_price_by_qty(double qty) {
    int64_t multd_qty = std::llround(qty * view_->qty_multiplier());
    int64_t acc_qty = 0;
    for (auto it = view_->BidBegin(); it != view_->BidEnd(); ++it) {
      acc_qty += it->qty;
      if (acc_qty >= multd_qty) {
        return it->price / view_->price_multiplier();
      }
    }
    return std::nullopt;
  }

 private:
  const view::BookBuilder* view_;
};  // class PyBookView

class PyTradeView {
 public:
  explicit PyTradeView(const view::TradeBuilder* view) : view_(view) {}
  ~PyTradeView() = default;

  const std::string& symbol() const { return view_->symbol(); }
  int64_t timestamp() const { return view_->timestamp(); }
  int64_t origin_timestamp() const { return view_->origin_timestamp(); }
  double price() const { return view_->price() / view_->price_multiplier(); }
  double qty() const { return view_->qty() / view_->qty_multiplier(); }
  int side() const { return view_->side(); }
  bool has_more_trade() const { return view_->has_more_trade(); }

  std::string to_str() {
    std::string side_str;
    switch (view_->side()) {
      case proto::TRADE_SIDE_BUY:
        side_str = "BUY";
        break;
      case proto::TRADE_SIDE_SELL:
        side_str = "SELL";
        break;
      default:
        side_str = "Unknown";
        break;
    }
    return fmt::format(
        "Trade(timestamp={}, symbol={}, side={}, price={}, qty={}, has_more_trade={})",
        timestamp(),
        symbol().c_str(),
        side_str.c_str(),
        price(),
        qty(),
        has_more_trade() ? "true" : "false");
  }

 private:
  const view::TradeBuilder* view_;
};  // class PyTradeView

class PyInstrumentView {
 public:
  PyInstrumentView(
      const view::StringJsonBuilder* view,
      const std::string& open_interest_column,
      const std::string& funding_rate_column,
      const std::string& mark_price_column,
      const std::string& indicative_funding_rate_column)
      : view_(view),
        open_interest_column_(open_interest_column),
        funding_rate_column_(funding_rate_column),
        mark_price_column_(mark_price_column),
        indicative_funding_rate_column_(indicative_funding_rate_column) {}
  ~PyInstrumentView() = default;

  int64_t timestamp() const { return view_->timestamp(); }

  double get_or_value(const std::string& column, double empty_value) const {
    auto jsonobj = view_->json_object();
    if (jsonobj.count(column) > 0) {
      return jsonobj[column].get<double>();
    } else {
      return empty_value;
    }
  }

  double open_interest() const {
    return get_or_value(open_interest_column_, std::numeric_limits<double>::quiet_NaN());
  }

  double funding_rate() const {
    return get_or_value(funding_rate_column_, std::numeric_limits<double>::quiet_NaN());
  }

  double mark_price() const {
    return get_or_value(mark_price_column_, std::numeric_limits<double>::quiet_NaN());
  }

  double indicative_funding_rate() const {
    return get_or_value(indicative_funding_rate_column_, std::numeric_limits<double>::quiet_NaN());
  }

  std::string to_str() { return view_->string(); }

 private:
  const view::StringJsonBuilder* view_;
  const std::string& open_interest_column_;
  const std::string& funding_rate_column_;
  const std::string& mark_price_column_;
  const std::string& indicative_funding_rate_column_;
};  // class PyInstrumentView

namespace {

void InvokeBookCallback(
    std::function<void(const PyBookView&)> callback,
    const view::BookBuilder* view) {
  PyBookView py_view(view);
  callback(py_view);
}

void InvokeTradeCallback(
    std::function<void(const PyTradeView&)> callback,
    const view::TradeBuilder* view) {
  PyTradeView py_view(view);
  callback(py_view);
}

void InvokeInstrumentCallback(
    std::function<void(const PyInstrumentView&)> callback,
    const std::string& open_interest_column,
    const std::string& funding_rate_column,
    const std::string& mark_price_column,
    const std::string& indicative_funding_rate_column,
    const view::StringJsonBuilder* view) {
  PyInstrumentView py_view(
      view,
      open_interest_column,
      funding_rate_column,
      mark_price_column,
      indicative_funding_rate_column);
  callback(py_view);
}

std::vector<filesystem::path> ConvertToPathVector(const std::vector<std::string>& strvec) {
  std::vector<filesystem::path> pathvec;
  for (auto&& path : strvec) {
    pathvec.emplace_back(path);
  }
  return pathvec;
}

}  // namespace

class PyProtoFeedBuilder {
 public:
  explicit PyProtoFeedBuilder(const std::string& sub_req) : holder_(new ProtoFeedBuilder(sub_req)) {
    builder_ = holder_.get();
  }
  explicit PyProtoFeedBuilder(FeedBuilderBase* builder) : builder_(builder) {}

  void update(std::string_view serialized) {
    CHECK_THROW(holder_, std::runtime_error);
    holder_->UpdateBySerializedProto(serialized);
  }

  void subscribe_book(const std::string& symbol, std::function<void(const PyBookView&)> callback) {
    builder_->SubscribeBook(
        symbol,
        std::bind(&InvokeBookCallback, callback, std::placeholders::_1));
  }

  void subscribe_trade(
      const std::string& symbol,
      std::function<void(const PyTradeView&)> callback) {
    builder_->SubscribeTrade(
        symbol,
        std::bind(&InvokeTradeCallback, callback, std::placeholders::_1));
  }

  void subscribe_instrument(
      const std::string& symbol,
      std::function<void(const PyInstrumentView&)> callback) {
    if (symbol.find(".Bitmex") != std::string::npos) {
      builder_->Subscribe<view::StringJsonBuilder>(
          proto::FEED_TYPE_CUSTOM,
          "bitmex_instrument",
          symbol,
          "json",
          std::bind(
              &InvokeInstrumentCallback,
              callback,
              "openInterest",
              "fundingRate",
              "markPrice",
              "indicativeFundingRate",
              std::placeholders::_1));
    } else if (
        symbol.find(".Okex") != std::string::npos &&
        builder_->sub_req().find("v3") != std::string::npos) {
      builder_->Subscribe<view::StringJsonBuilder>(
          proto::FEED_TYPE_CUSTOM,
          "okex_instrument",
          symbol,
          "json",
          std::bind(
              &InvokeInstrumentCallback,
              callback,
              "open_interest",
              "funding_rate",
              "mark_price",
              "indicative_funding_rate",
              std::placeholders::_1));
    }
  }

  void* cc_builder() { return builder_; }

 private:
  FeedBuilderBase* builder_;
  std::unique_ptr<ProtoFeedBuilder> holder_;
};  // class PyProtoFeedBuilder

class PyFeedCacheRunner {
 public:
  PyFeedCacheRunner(const std::string& date_str, const std::vector<std::string>& feed_cache_dir)
      : runner_(
            std::bind(
                &PyFeedCacheRunner::OnFeedBuilderReset,
                this,
                std::placeholders::_1,
                std::placeholders::_2),
            date_str,
            ConvertToPathVector(feed_cache_dir)) {}

  void set_book_reset_callback(
      std::function<void(const FeedCacheKey&, PyProtoFeedBuilder*)> callback) {
    book_reset_callback_ = callback;
  }

  void add_cache_key(
      const std::string& machine,
      const std::string& worker_id,
      const std::string& sub_req,
      const std::vector<std::string>& symbols) {
    runner_.AddCacheKey(FeedCacheKey {machine, worker_id, sub_req}, symbols);
  }

  void add_cache_key_with_feed_arb(
      const std::string& machine,
      const std::string& sub_req,
      const std::vector<std::string>& symbols) {
    runner_.AddCacheKeyWithFeedArb(machine, sub_req, symbols);
  }

  void run(std::optional<int64_t> from_timestamp, std::optional<int64_t> to_timestamp) {
    runner_.Run(from_timestamp, to_timestamp);
  }

 private:
  void OnFeedBuilderReset(const FeedCacheKey& feed_cache_key, FeedBuilderBase* feed_builder) {
    if (book_reset_callback_) {
      PyProtoFeedBuilder py_feed_builder(feed_builder);
      book_reset_callback_(feed_cache_key, &py_feed_builder);
    }
  }

  std::function<void(const FeedCacheKey&, PyProtoFeedBuilder*)> book_reset_callback_;
  FeedCacheRunner runner_;
};  // class PyFeedCacheRunner

PYBIND11_MODULE(pyfastfeed, m) {
  py::class_<PyBookLevel>(m, "_BookLevel")
      .def_readonly("price", &PyBookLevel::price)
      .def_readonly("qty", &PyBookLevel::qty)
      .def_readonly("count", &PyBookLevel::count)
      .def("__repr__", &PyBookLevel::to_str);

  py::class_<PyBookIterator>(m, "_BookIterator")
      .def("__iter__", [](PyBookIterator& it) -> PyBookIterator& { return it; })
      .def("__next__", &PyBookIterator::next);

  py::class_<PyBookView>(m, "_BookView")
      .def_property_readonly("symbol", &PyBookView::symbol)
      .def_property_readonly("timestamp", &PyBookView::timestamp)
      .def_property_readonly("exchange_timestamp", &PyBookView::origin_timestamp)
      .def("has_ask", &PyBookView::has_ask)
      .def("has_bid", &PyBookView::has_bid)
      .def("depth_ask", &PyBookView::depth_ask)
      .def("depth_bid", &PyBookView::depth_bid)
      .def("ask0", &PyBookView::ask0)
      .def("bid0", &PyBookView::bid0)
      .def("get_ask_array", &PyBookView::get_ask_array)
      .def("get_bid_array", &PyBookView::get_bid_array)
      .def("iter_ask", &PyBookView::iter_ask)
      .def("iter_bid", &PyBookView::iter_bid)
      .def("get_accumulated_ask_qty", &PyBookView::get_accumulated_ask_qty)
      .def("get_accumulated_bid_qty", &PyBookView::get_accumulated_bid_qty)
      .def("get_ask_price_by_amount", &PyBookView::get_ask_price_by_amount)
      .def("get_bid_price_by_amount", &PyBookView::get_bid_price_by_amount)
      .def("get_ask_price_by_qty", &PyBookView::get_ask_price_by_qty)
      .def("get_bid_price_by_qty", &PyBookView::get_bid_price_by_qty);

  py::class_<PyTradeView>(m, "_TradeView")
      .def_property_readonly("symbol", &PyTradeView::symbol)
      .def_property_readonly("timestamp", &PyTradeView::timestamp)
      .def_property_readonly("exchange_timestamp", &PyTradeView::origin_timestamp)
      .def_property_readonly("price", &PyTradeView::price)
      .def_property_readonly("qty", &PyTradeView::qty)
      .def_property_readonly("side", &PyTradeView::side)
      .def_property_readonly("has_more_trade", &PyTradeView::has_more_trade)
      .def("__repr__", &PyTradeView::to_str);

  py::class_<PyInstrumentView>(m, "_InstrumentView")
      .def_property_readonly("timestamp", &PyInstrumentView::timestamp)
      .def_property_readonly("open_interest", &PyInstrumentView::open_interest)
      .def_property_readonly("funding_rate", &PyInstrumentView::funding_rate)
      .def_property_readonly("mark_price", &PyInstrumentView::mark_price)
      .def_property_readonly("indicative_funding_rate", &PyInstrumentView::indicative_funding_rate)
      .def("__repr__", &PyInstrumentView::to_str);

  py::class_<PyProtoFeedBuilder>(m, "_ProtoFeedBuilder")
      .def_property_readonly("cc_builder", &PyProtoFeedBuilder::cc_builder)
      .def(py::init<const std::string&>())
      .def("update", &PyProtoFeedBuilder::update)
      .def("subscribe_book", &PyProtoFeedBuilder::subscribe_book)
      .def("subscribe_trade", &PyProtoFeedBuilder::subscribe_trade)
      .def("subscribe_instrument", &PyProtoFeedBuilder::subscribe_instrument);

  py::class_<FeedCacheKey>(m, "_FeedCacheKey")
      .def(py::init())
      .def_readwrite("machine", &FeedCacheKey::machine)
      .def_readwrite("worker_id", &FeedCacheKey::worker_id)
      .def_readwrite("sub_req", &FeedCacheKey::sub_req)
      .def("__repr__", &FeedCacheKey::ToString);

  py::class_<PyFeedCacheRunner>(m, "_FeedCacheRunner")
      .def(py::init<const std::string&, const std::vector<std::string>&>())
      .def("set_book_reset_callback", &PyFeedCacheRunner::set_book_reset_callback)
      .def("add_cache_key", &PyFeedCacheRunner::add_cache_key)
      .def("add_cache_key_with_feed_arb", &PyFeedCacheRunner::add_cache_key_with_feed_arb)
      .def("run", &PyFeedCacheRunner::run);
}

}  // namespace fastfeed
